#include <cuda.h>
#include <cuda_runtime.h>
#include <stdio.h>

// 错误检查宏
#define CUDA_CHECK_ERROR(err)                                                           \
    if (err != cudaSuccess)                                                             \
    {                                                                                   \
        fprintf(stderr, "CUDA Error: %s (err_num=%d)\n", cudaGetErrorString(err), err); \
        exit(EXIT_FAILURE);                                                             \
    }

// 首次归约内核：计算每个块的部分和
__global__ void reduce_sum_of_squares(const float *inp, float *partial_sums, int nr_elems)
{
    extern __shared__ float sdata[];

    unsigned int tid = threadIdx.x;
    unsigned int idx = blockIdx.x * blockDim.x * 2 + threadIdx.x;

    float sum = 0.0f;

    // 每个线程加载两个元素（如果存在）
    if (idx < nr_elems)
        sum += inp[idx] * inp[idx];
    if (idx + blockDim.x < nr_elems)
        sum += inp[idx + blockDim.x] * inp[idx + blockDim.x];

    sdata[tid] = sum;
    __syncthreads();

    // 归约：逐步减少到一个元素
    for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1)
    {
        if (tid < s)
        {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }

    // 将每个块的部分和写入全局内存
    if (tid == 0)
    {
        partial_sums[blockIdx.x] = sdata[0];
    }
}

// 最终归约内核：将部分和进一步归约到一个值
__global__ void final_reduce_sum(float *partial_sums, float *out, int nr_partials)
{
    extern __shared__ float sdata[];

    unsigned int tid = threadIdx.x;
    unsigned int idx = threadIdx.x;

    float sum = 0.0f;
    if (idx < nr_partials)
        sum = partial_sums[idx];

    sdata[tid] = sum;
    __syncthreads();

    // 归约
    for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1)
    {
        if (tid < s)
        {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }

    // 写入最终结果
    if (tid == 0)
    {
        *out = sdata[0];
    }
}

// 主机函数：计算欧几里得范数
float compute_euclidean_norm(const float *d_inp, int nr_elems)
{
    float *d_partial_sums;
    int threads = 256;
    int blocks = (nr_elems + threads * 2 - 1) / (threads * 2);

    // 分配部分和的设备内存
    CUDA_CHECK_ERROR(cudaMalloc((void **)&d_partial_sums, blocks * sizeof(float)));

    // 首次归约
    reduce_sum_of_squares<<<blocks, threads, threads * sizeof(float)>>>(d_inp, d_partial_sums, nr_elems);
    CUDA_CHECK_ERROR(cudaGetLastError());

    // 如果部分和的数量较大，可能需要多次归约
    int current_blocks = blocks;
    while (current_blocks > 1)
    {
        int threads_final = 256;
        int blocks_final = (current_blocks + threads_final * 2 - 1) / (threads_final * 2);

        // 重新分配部分和的内存（覆盖原有部分和）
        float *d_new_partial_sums;
        CUDA_CHECK_ERROR(cudaMalloc((void **)&d_new_partial_sums, blocks_final * sizeof(float)));

        // 进行一次归约
        reduce_sum_of_squares<<<blocks_final, threads_final, threads_final * sizeof(float)>>>(d_partial_sums, d_new_partial_sums, current_blocks);
        CUDA_CHECK_ERROR(cudaGetLastError());

        // 释放旧的部分和内存
        CUDA_CHECK_ERROR(cudaFree(d_partial_sums));

        // 更新部分和指针和块数
        d_partial_sums = d_new_partial_sums;
        current_blocks = blocks_final;
    }

    // 最终归约到一个值
    float h_out;
    final_reduce_sum<<<1, threads, threads * sizeof(float)>>>(d_partial_sums, &h_out, current_blocks);
    CUDA_CHECK_ERROR(cudaGetLastError());

    // 释放部分和的设备内存
    CUDA_CHECK_ERROR(cudaFree(d_partial_sums));

    // 计算平方根，得到欧几里得范数
    return sqrtf(h_out);
}

int main()
{
    // 示例输入
    int nr_elems = 1 << 20; // 例如 1M 元素
    size_t size = nr_elems * sizeof(float);

    // 分配主机内存并初始化
    float *h_inp = (float *)malloc(size);
    for (int i = 0; i < nr_elems; i++)
    {
        h_inp[i] = 1.0f; // 示例数据，范数应为 sqrt(nr_elems)
    }

    // 分配设备内存
    float *d_inp;
    CUDA_CHECK_ERROR(cudaMalloc((void **)&d_inp, size));

    // 将数据从主机复制到设备
    CUDA_CHECK_ERROR(cudaMemcpy(d_inp, h_inp, size, cudaMemcpyHostToDevice));

    // 计算欧几里得范数
    float norm = compute_euclidean_norm(d_inp, nr_elems);

    printf("Euclidean Norm: %f\n", norm); // 对于全1数组，应为 sqrt(nr_elems)

    // 清理
    CUDA_CHECK_ERROR(cudaFree(d_inp));
    free(h_inp);

    return 0;
}
